"""
构造行列式
"""

import numpy
import torch
from torch import nn



class DetLayer1(nn.Module):
    def __init__(self, nup, ndn):
        ''''''
        super().__init__()
        self.nup = nup
        self.ndn = ndn
        self.nel = self.nup + self.ndn
        #传过来的参数是（nup，4）和（ndn，4）
        #最后要产生一个(nup, nup)的矩阵并求det，每一个元素都是dot((4),(4))得到的
        self.Wup = nn.Parameter(torch.randn(self.nup, 4))
        self.gup = nn.Parameter(torch.zeros(self.nup))
        #
        self.Wdn = nn.Parameter(torch.randn(self.ndn, 4))
        self.gdn = nn.Parameter(torch.zeros(self.ndn))
        #\Sigma函数，只有核心有一个原子(m=1)
        self.decay_up = lambda t: torch.linalg.vector_norm(t, dim=1, keepdim=True)
        self.pi_up_ai = nn.Parameter(torch.ones(self.nup, 1))
        #
        self.decay_dn = lambda t: torch.linalg.vector_norm(t, dim=1, keepdim=True)
        self.pi_dn_ai = nn.Parameter(torch.ones(self.ndn, 1))


    def forward(
            self,
            rupdn: torch.Tensor,
            prev_h_ai: torch.Tensor,
            prev_h_aibj: torch.Tensor
        ):
        #
        rup = rupdn[:self.nup, :]
        self.decay_up_mask = torch.exp(-self.decay_up(rup))
        self.decay_up_mask = self.pi_up_ai * self.decay_up_mask
        rdn = rupdn[self.nup:, :]
        self.decay_dn_mask = torch.exp(-self.decay_dn(rdn))
        self.decay_dn_mask = self.pi_dn_ai * self.decay_dn_mask
        #构造行列式
        #输入的粒子位置交换，到现在的时候交换的是列指标，所以给每行加或者乘不同的数字没有关系
        #只要行内的操作相同，列的交换就能保持
        mat_up = torch.matmul(self.Wup, prev_h_ai[:self.nup, :].T)
        mat_up = mat_up + self.gup.unsqueeze(dim=1)
        mat_up = self.decay_up_mask * mat_up
        #
        mat_dn = torch.matmul(self.Wdn, prev_h_ai[self.nup:, :].T)
        mat_dn = mat_dn + self.gdn.unsqueeze(dim=1)
        mat_dn = self.decay_dn_mask * mat_dn
        #
        #print("matup", mat_up)
        #print("matdn", mat_dn)
        _, psiup = torch.linalg.slogdet(mat_up)
        _, psidn = torch.linalg.slogdet(mat_dn)
        #print(s1, s2)
        return psiup+psidn
